#!/usr/bin/python
import csv
import h5py
import math
from numpy import genfromtxt
import os
import sys
import sys, struct
import vtk
#from vtk.vtkGraphicsPython import vtkPlaneSource, vtkStreamTracer
from vtk.util import numpy_support as nsup

import math as pm
import matplotlib.pyplot as plt
import numpy as np
import pylab as p


# This is statement is required by the build system to query build info
if __name__ == '__build__':
    raise Exception

'''
    Author: Brad Hollister.
    Started: March 20th, 2015
    Code advects streamlines for each ensemble member and increments the number of times a streamline crosses a given cell.
'''

#from vtk import vtkAlgorithm as vta

PREFIX_U = 'u_r_'
PREFIX_V = 'v_r_'
PREFIX_RHO = 'rho_r_'
PATH = '/home/brad/DATA/ts00050/'  
#PATH = '/home/behollis/DATA/in/ts00050/'  
OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'
#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'

'''
Lock-exchange data parameters.

> Nx: 128
> Ny: 128
> T: 15  (non-dimensional)
> dt: 4.0000e-04
> nu: 2.5000e-04
> kappa: 2.5000e-04
> S: 15 or 20  (total number of modes)
> PlotIntrvl: 50  (number of time-step/interval between plot times)
> var: 0.0500
> MC: 1000
> PrRa: 1
> The realizations are sample path evolution of the flow: I believe that
> there are 1000 of them in what we give you.
'''

def readMember(mem):
    ''' Read numpy matrix from disk and return tuple of realization data. '''
    filename_u = PREFIX_U + str(mem) + '.txt'
    filename_v = PREFIX_V + str(mem) + '.txt'
    filename_rho = PREFIX_RHO + str(mem) + '.txt'
    
    u = genfromtxt(PATH + filename_u)
    u = -1.0*u[0:X,0:Y]
    
    v = genfromtxt(PATH + filename_v)
    v = v[0:X,0:Y]
    #rho = genfromtxt(PATH + filename_rho)
    
    return (u,v)

def writeElapsedTime(start, end):
    import os
    systime = os.times()
    tout = open(OUTPUT_DATA_DIR+'lock.exchange.timing.{0}.txt'.format(systime[4]), 'w') #sys time
    tout.write('Elapsed time: {0}'.format(end - start))
    tout.close()
    

if __name__ == "__main__":  
    
    import random,time
    start = time.time()
    
    TOTMEMS = 30
    X = 125
    Y = 125
    DIM = 2
        
    members = set()
    while len(members) < TOTMEMS:
        members.add(random.randint(1,1000))
   
    output = h5py.File( OUTPUT_DATA_DIR+'lockSlines.1100steps.hdf5', 'w')
   
    st = vtk.vtkStreamTracer()
    
    st.SetMaximumPropagation(1000)
    st.SetMaximumNumberOfSteps(1100)
    st.SetInitialIntegrationStep(0.01)
    st.SetMinimumIntegrationStep(0.1)
    st.SetMaximumIntegrationStep(0.1)
    st.SetTerminalSpeed(0.00001)
    st.SetIntegratorTypeToRungeKutta45()
    st.SetComputeVorticity(False)
    
    #generate streamlines for each slice of the corresponding member vector field
    idx = 1
    for mem in members:
        print 'computing member = {0}, {1} of {2}'.format(mem,idx,TOTMEMS)
        #increment member count
        idx += 1
        
        u, v = readMember(mem)
        npts = u.shape[0] * u.shape[1]

        Vel = np.zeros((npts,3))
        Vel[:,0] = u.reshape(npts,)
        Vel[:,1] = v[:,:].reshape(npts,)

        velArray = nsup.numpy_to_vtk(Vel)
        velArray.SetName("velocity")

        image = vtk.vtkImageData()
        image.SetDimensions(X, Y, 1)
        image.GetPointData().AddArray(velArray)
        
        st.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "velocity")
        st.SetInputData(image)
        
        for y in range(0,Y):
            for x in range(0,X):
                
                #print 'writing sline: x={0} y={1}'.format(x,y)
                
                st.SetStartPosition(x,y,0)
                st.SetIntegrationDirectionToForward()
                st.Update()
                
                vtkPolyData_vtp = st.GetOutput()#w.GetInput()
                #print 'cells: ' + str(vtkPolyData_vtp.GetNumberOfCells())
                cells = vtkPolyData_vtp.GetNumberOfCells()
            
                dir = '/mem' + str(mem).zfill(4) + '/x' + str(x).zfill(3) + '/y' + str(y).zfill(3) 
                
                vtkPolyLine_slF = vtkPolyData_vtp.GetCell( 0 )
                vtkPoints_ptsF = vtkPolyLine_slF.GetPoints()
                
               
                #copy streamline points to numpy array
                num_ptsF = vtkPoints_ptsF.GetNumberOfPoints()
            
                slnpF = np.ndarray(shape=(2,num_ptsF))
                
                #print 'forward points: ' + str(num_ptsF)
                
                for pt_id in range(0,num_ptsF):
                    pt_tupleF = vtkPoints_ptsF.GetPoint(pt_id)
                    
                    slnpF[0][pt_id] = pt_tupleF[0]
                    slnpF[1][pt_id] = pt_tupleF[1]      
            
                st.SetStartPosition(x,y,0)
                st.SetIntegrationDirectionToBackward()
                st.Update()
                
                vtkPolyData_vtp = st.GetOutput()#w.GetInput()
                #print 'cells: ' + str(vtkPolyData_vtp.GetNumberOfCells())
                cells = vtkPolyData_vtp.GetNumberOfCells()
            
                vtkPolyLine_slB = vtkPolyData_vtp.GetCell( 0 )
                vtkPoints_ptsB = vtkPolyLine_slB.GetPoints()
                
                #copy streamline points to numpy array
                num_ptsB = vtkPoints_ptsB.GetNumberOfPoints()
                slnpB = np.ndarray(shape=(2,num_ptsB))
                
                #print 'backward points: ' + str(num_ptsB)
                
                for pt_id in range(0, num_ptsB):
                    pt_tupleB = vtkPoints_ptsB.GetPoint(pt_id)
                
                    slnpB[0][pt_id] = pt_tupleB[0]
                    slnpB[1][pt_id] = pt_tupleB[1] 
                     
                slnpB = np.fliplr(slnpB)   
                allpts = np.concatenate((slnpB[:,1:], slnpF), axis=1)
                
                output[dir] = allpts
            
    print "finished!"
    
    end = time.time()
    writeElapsedTime(start, end)
        
    output.close()
